import torch
from torch import nn

class HeatmapFocalLoss(nn.Module):
    def __init__(self, weight=1.0):
        super(HeatmapFocalLoss, self).__init__()
        
    def forward(self, pred, target, mask):
        pos_inds = target.eq(1).float()
        neg_inds = target.lt(1).float()
        neg_weights = torch.pow(1 - target, 4)
        pos_loss = torch.log(pred) * torch.pow(1 - pred, 2) * pos_inds
        neg_loss = torch.log(1 - pred) * torch.pow(pred, 2) * neg_weights * neg_inds

        # normalize
        num_pos = torch.clamp(torch.sum(pos_inds), min=1, max=1e30)
        if mask is not None:
            pos_loss = (pos_loss * mask).sum()
            neg_loss = (neg_loss * mask).sum()
        else:
            pos_loss = pos_loss.sum()
            neg_loss = neg_loss.sum()
        return -(pos_loss + neg_loss) / num_pos


